package com.lzx.ai.utils;

/**
 * @author 程序员星星
 * @date 2023/4/14
 * @Description
 */

import edu.stanford.nlp.ling.CoreAnnotations;
import edu.stanford.nlp.ling.CoreLabel;
import edu.stanford.nlp.pipeline.Annotation;
import edu.stanford.nlp.pipeline.StanfordCoreNLP;
import edu.stanford.nlp.util.CoreMap;

import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

public class KeywordExtractor {

    private static final String DEFAULT_MODEL_PATH = "";

    private StanfordCoreNLP pipeline;

    public KeywordExtractor() {
        Properties props = new Properties();
        props.setProperty("annotators", "tokenize, ssplit, pos, lemma");
        pipeline = new StanfordCoreNLP(props);
    }

    public List<String> extract(String text) {
        Annotation document = new Annotation(text);
        pipeline.annotate(document);
        List<CoreMap> sentences = document.get(CoreAnnotations.SentencesAnnotation.class);
        List<String> keywords = new ArrayList<>();
        for (CoreMap sentence : sentences) {
            for (CoreLabel token : sentence.get(CoreAnnotations.TokensAnnotation.class)) {
                String word = token.get(CoreAnnotations.TextAnnotation.class);
                String pos = token.get(CoreAnnotations.PartOfSpeechAnnotation.class);
                if (isKeyword(pos)) {
                    keywords.add(word);
                }
            }
        }
        return keywords;
    }

    private boolean isKeyword(String pos) {
        return pos.equals("NN") || pos.equals("NNS") || pos.equals("NNP") || pos.equals("NNPS");
    }
}